"""
BEL (Binary Error Locating) Engine

基于 URSA-8B 推理框架实现的二元错误定位引擎，
用于自动为错误推理路径逐步打标签，生成 S_BEL 过程监督数据。

严格按照设计文档要求，复用现有代码，实现同步执行的两阶段流程：
1. 生成阶段：generator.py 生成 rollouts.jsonl
2. 标注阶段：labeler.py 标注 S_BEL.jsonl

优化版本：
- 移除重复代码，最大化复用现有函数
- 简化自适应采样，使用固定配置
- 统一图片路径字段处理
- 修复二分查找算法逻辑
- 增强类型安全和输入验证

BEL（二元错误定位）引擎模块

在模块导入时自动注册 URSA 模型到 transformers，确保 vLLM 能够正确识别
"""

import sys
import os

# 确保能够导入 models 模块
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

try:
    # 导入并注册 URSA 模型
    from transformers import CONFIG_MAPPING, AutoConfig, AutoModelForCausalLM
    from models.ursa_model.configuration_ursa import UrsaConfig
    from models.ursa_model.modeling_ursa import UrsaForConditionalGeneration
    
    # 注册 URSA 模型配置和模型类
    if "ursa" not in CONFIG_MAPPING:
        CONFIG_MAPPING.register("ursa", UrsaConfig)
        AutoConfig.register("ursa", UrsaConfig)
        AutoModelForCausalLM.register(UrsaConfig, UrsaForConditionalGeneration)
        print("✅ 成功注册 URSA 模型到 transformers")
    else:
        print("ℹ️  URSA 模型已经注册")
        
except ImportError as e:
    print(f"⚠️  无法注册 URSA 模型: {e}")
    print("请确保在项目根目录运行，并已安装相关依赖")

__version__ = "0.3.1"

from typing import Dict, List, Optional, Any, Union

from .engine import BinaryErrorLocator
from .utils import (
    parse_trace, split_steps, format_steps_with_labels, insert_labels,
    load_config, validate_config, setup_logging, ProgressBar, count_lines,
    validate_answer, normalize_image_path, create_llm_instance, create_sampling_params,
    Trace, validate_bel_format
)

# 导出主要模块
from .generator import RolloutGenerator
from .labeler import BinaryErrorLabeler

__all__ = [
    "BinaryErrorLocator",
    "parse_trace",
    "split_steps",
    "format_steps_with_labels",
    "insert_labels",
    "load_config",
    "validate_config", 
    "setup_logging",
    "ProgressBar",
    "count_lines",
    "validate_answer",
    "normalize_image_path",
    "create_llm_instance",
    "create_sampling_params",
    "Trace",
    "validate_bel_format",
    "RolloutGenerator",
    "BinaryErrorLabeler"
] 